import numpy as np
import torch
import torch.nn as nn


class PytorchInit:
    def __init__(self):
        pass
        
    def init_weight(self, layer):
        if type(layer) == nn.Conv2d:
            nn.init.constant_(layer.weight, self.constant)
            nn.init.constant_(layer.bias, 0)
        elif type(layer) == nn.Conv3d:
            nn.init.constant_(layer.weight, self.constant)
            nn.init.constant_(layer.bias, 0)
        elif type(layer) == nn.Conv1d:
            nn.init.constant_(layer.weight, self.constant)
            nn.init.constant_(layer.bias, 0)
        elif type(layer) == nn.Linear:
            nn.init.constant_(layer.weight, self.constant)
            nn.init.constant_(layer.bias, 0)
        elif type(layer) == nn.LSTM:
            for name, param in layer.named_parameters():
                if name.startswith("weight"):
                    nn.init.constant_(param, self.constant)
                else:
                    nn.init.zeros_(param, 0)
                    
    # # another init fun
    # def net_init(self, net, constant):
    #     self.constant = constant
    #     net = net.apply(self.init_weight)
    #     return net
    
    
    def net_init(self, net, constant):
        stat_dict = net.state_dict()
        param_dic = {}
        for k in stat_dict:
            shape = stat_dict[k].shape
            if "bias" in k:
                v = nn.init.constant_(torch.empty(shape),0)
            else:
                v = nn.init.constant_(torch.empty(shape),constant)
            stat_dict[k] = v
        net.load_state_dict(stat_dict)
        return net
